
%% This function for epoching the walking data

% Input:
% EEG is the preprocessed data
% evenOnsetName is the name of event onset; eventOnsetName = 'UserStartWalking'
% conditionsName is the list of walking condition; conditionsName = {'Walk12' 'Walk22' 'Walk13' 'Walk23'}
% epochLength is length of epoch; lengthWalkEpoch = [-2.5 12];
% epochDIr is the directory for the generated epoch data

function epoch(EEG, eventOnsetName, conditionsName, lengthWalkEpoch, epochDir)

disp(['EPOCHING DATASET: ' EEG.filename]);
    
eeglab_complete_path = which('eeglab');
eeglab_split_path = strsplit(eeglab_complete_path,'eeglab.m');
eeglabDir = eeglab_split_path{1};
   
%% EPOCH THE DATA INTO WALK CONDIOTIONS
EEG = pop_epoch( EEG, {  eventOnsetName  }, lengthWalkEpoch, 'epochinfo', 'yes');
EEG = pop_rmbase( EEG, [lengthWalkEpoch(1) 0]*1000);    % Change from seconds to miliseconds
EEG = eeg_checkset(EEG);

%% Segment the data for further analysis into walk 12, walk 13, walk 22, walk23
% Information to segment walking data
% Walk collide sphere sequency:
% (walk1..)	(walk2..)
% 1 2 3		4 5
% 6 7 		8 9 10
% 11 12		13 14
% 15 16 17	18 19 20
epochIdx12 = [];
epochIdx13 = [];
epochIdx22 = [];
epochIdx23 = [];

for epochCur =1:size(EEG.epoch,2)
    tempEventType = EEG.epoch(epochCur).eventtype;
    for eventCur = 1:size(tempEventType,2)
%         walk12
        if strcmp(tempEventType{eventCur},'CollideWithSphere6')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere7')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere11')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere12')
            epochIdx12 = [epochIdx12; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         walk13        
        if strcmp(tempEventType{eventCur},'CollideWithSphere1')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere2')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere3')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere15')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere16')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere17')
            epochIdx13 = [epochIdx13; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         walk22
        if strcmp(tempEventType{eventCur},'CollideWithSphere4')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere5')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere13')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere14')
            epochIdx22 = [epochIdx22; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         walk23        
        if strcmp(tempEventType{eventCur},'CollideWithSphere8')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere9')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere10')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere18')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere19')||...
                strcmp(tempEventType{eventCur},'CollideWithSphere20')
            epochIdx23 = [epochIdx23; epochCur];   
            break; % CAREFUL MISSING DATA
        end        
    end
end

% Set the lable to segment the data
trialrej12 = ones(1,size(EEG.epoch,2));
trialrej12(epochIdx12) = 0;

trialrej13 = ones(1,size(EEG.epoch,2));
trialrej13(epochIdx13) = 0;

trialrej22 = ones(1,size(EEG.epoch,2));
trialrej22(epochIdx22) = 0;

trialrej23 = ones(1,size(EEG.epoch,2));
trialrej23(epochIdx23) = 0;

eeglab redraw
EEG = eeg_checkset(EEG);


%%
mkdir([epochDir '\' EEG.filename(1:3)]);

EEG12 = pop_rejepoch( EEG, trialrej12, 0);
EEG22 = pop_rejepoch( EEG, trialrej22, 0);
EEG13 = pop_rejepoch( EEG, trialrej13, 0);
EEG23 = pop_rejepoch( EEG, trialrej23, 0);

% automatic artifact rejection
EEG12 = pop_autorej(EEG12, 'nogui','on','threshold',1000,'eegplot','off');
EEG13 = pop_autorej(EEG13, 'nogui','on','threshold',1000,'eegplot','off');
EEG22 = pop_autorej(EEG22, 'nogui','on','threshold',1000,'eegplot','off');
EEG23 = pop_autorej(EEG23, 'nogui','on','threshold',1000,'eegplot','off');

[EEG12] = eeg_checkset(EEG12);
[EEG22] = eeg_checkset(EEG22);
[EEG13] = eeg_checkset(EEG13);
[EEG23] = eeg_checkset(EEG23);

%% check the epoch data quality here, will be used it later
[EEG12, timeWarp12]  = thong_timeWarpInformation(EEG12, 2);
[EEG22, timeWarp22]  = thong_timeWarpInformation(EEG22, 2);
[EEG13, timeWarp13]  = thong_timeWarpInformation(EEG13, 3);
[EEG23, timeWarp23]  = thong_timeWarpInformation(EEG23, 3);

%% Saving the data

[EEG12] = eeg_checkset(EEG12);
[EEG22] = eeg_checkset(EEG22);
[EEG13] = eeg_checkset(EEG13);
[EEG23] = eeg_checkset(EEG23);

pop_saveset( EEG12, 'filename', [EEG.filename(1:3) '_' conditionsName{1} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG22, 'filename', [EEG.filename(1:3) '_' conditionsName{2} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG13, 'filename', [EEG.filename(1:3) '_' conditionsName{3} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG23, 'filename', [EEG.filename(1:3) '_' conditionsName{4} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);

disp([EEG12.filename ' has ' num2str(size(EEG12.epoch,2)) ' epochs']); 
disp([EEG22.filename ' has ' num2str(size(EEG22.epoch,2)) ' epochs']); 
disp([EEG13.filename ' has ' num2str(size(EEG13.epoch,2)) ' epochs']); 
disp([EEG23.filename ' has ' num2str(size(EEG23.epoch,2)) ' epochs']); 

end

